package aggregate

import (
	"context"
	"go.uber.org/ratelimit"
	"nprobe/core"
	"nprobe/utils/routines"
	"time"
)

type HeapMapAggregator struct {
	Ctx         context.Context
	Name        string
	Task        *core.ProbeTask
	FirstSpanID int64
	Bucket      *core.PktRecordBucket
}

func NewHeapMapAggregator(
	ctx context.Context,
	name string,
	task *core.ProbeTask,
	bucket *core.PktRecordBucket,
) *HeapMapAggregator {
	return &HeapMapAggregator{
		Ctx:    ctx,
		Name:   name + "-aggregator",
		Task:   task,
		Bucket: bucket,
	}
}

func (a *HeapMapAggregator) StartAggregator() error {
	routines.GoWithRecover(a.Name, a.doAggregate)
	return nil
}

func (a *HeapMapAggregator) StopAggregator() error {
	return nil
}

func (a *HeapMapAggregator) doAggregate() {
	limiter := ratelimit.New(1, ratelimit.Per(time.Duration(a.Task.SpanMs)*time.Millisecond))
	for {
		select {
		case <-a.Ctx.Done():
			a.StopAggregator()
			return
		default:
			limiter.Take()
			a.tryAggregate()
		}
	}
}

func (a *HeapMapAggregator) tryAggregate() {
	ts := time.Now().UnixNano()
	timeoutSpanID := (ts-a.Task.TimeoutMs*int64(time.Millisecond))/(a.Task.SpanMs*int64(time.Millisecond)) - 1
	var spanRecord *core.SpanRecord
	for {
		pktRecord := a.Bucket.Peek()
		if pktRecord == nil || pktRecord.SpanID > timeoutSpanID {
			// 收集完，统计打印
			aggregateSpanRecordAndPrint(a.FirstSpanID, spanRecord, a.Task.Role)
			break
		}

		_ = a.Bucket.Pop()
		if spanRecord == nil || pktRecord.SpanID > spanRecord.SpanID {
			// 一个新的已经超时的span
			if a.FirstSpanID == 0 && pktRecord.SpanID != 0 {
				// 第一个span
				a.FirstSpanID = pktRecord.SpanID
			}
			aggregateSpanRecordAndPrint(a.FirstSpanID, spanRecord, a.Task.Role)
			spanRecord = core.NewSpanRecord(pktRecord.SpanID)
		}
		spanRecord.PktRecordMap[pktRecord.Key] = pktRecord
	}
}
